import os
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import List, Optional, Any, Dict
import subprocess

@dataclass
class DLTask:
    command: str
    task_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    created_at: datetime = field(default_factory=datetime.now)
    num_gpu: int = 0  # 默认不需要GPU
    mem_cpu: Optional[float] = None  # 默认内存占用为0MB
    mem_per_gpu: Optional[float] = None
    require_devices: Optional[List[str]] = None
    allocated_devices: List[str] = field(default_factory=list)
    process: Optional[subprocess.Popen] = None
    status: str = "PENDING"
    max_retries: int = 3
    retries: int = 0
    last_failure: Optional[str] = None
    env: Dict[str, str] = field(default_factory=lambda: os.environ.copy()) 
    pwd: str = field(default_factory="~/")

    def __post_init__(self):
        unsafe_commands = {"rm", "kill", "shutdown"}
        if any(cmd in self.command for cmd in unsafe_commands):
            raise ValueError(f"命令包含危险操作 '{self.command}'")

    def metadata(self) -> str:
        metadata_items = [
        f"command: {self.command}",
        f"created_at: {self.created_at}",
        f"require_devices: {self.require_devices}",
        f"allocated_devices: {self.allocated_devices}",
        f"gpus: {self.num_gpu}",
        f"cpu_mem: {self.mem_cpu}",
        f"gpu_mem: {self.mem_per_gpu}",
        f"retries: {self.retries}",
        f"env: {self.env}",
        f"pwd: {self.pwd}",
        "" # 空行
        ]
        return "\n".join(metadata_items)